#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <malloc.h>
#include <sys/time.h>
#include "aceMesh_runtime_c.h"
// double A[N/TS][M/TS][TS][TS];
// double B[M/TS][P/TS][TS][TS];
// double C[N/TS][P/TS][TS][TS];
double ****A,****B,****C;
void initialize(int N,int M,int P,int TS) 
{
   int i,j,ii,jj;
   A=(double****)malloc(N/TS*sizeof(double***));
   B=(double****)malloc(M/TS*sizeof(double***));
   C=(double****)malloc(N/TS*sizeof(double***));
   for(i=0;i<N/TS;i++){
      A[i]=(double***)malloc(M/TS*sizeof(double**));
      C[i]=(double***)malloc(P/TS*sizeof(double**));
      for(j=0;j<M/TS;j++){
         A[i][j]=(double**)malloc(TS*sizeof(double*));
         for(ii=0;ii<TS;ii++){
            A[i][j][ii]=(double*)malloc(TS*sizeof(double));
            for(jj=0;jj<TS;jj++){
               A[i][j][ii][jj]=1.0;
            }
         }
      }
      for(j=0;j<P/TS;j++){
         C[i][j]=(double**)malloc(TS*sizeof(double*));
         for(ii=0;ii<TS;ii++){
            C[i][j][ii]=(double*)malloc(TS*sizeof(double));
            for(jj=0;jj<TS;jj++){
               C[i][j][ii][jj]=3.0;
            }
         }
      }
   }
   for(i=0;i<M/TS;i++){
      B[i]=(double***)malloc(P/TS*sizeof(double**));
      for(j=0;j<P/TS;j++){
         B[i][j]=(double**)malloc(TS*sizeof(double*));
         for(ii=0;ii<TS;ii++){
            B[i][j][ii]=(double*)malloc(TS*sizeof(double));
            for(jj=0;jj<TS;jj++){
               B[i][j][ii][jj]=2.0;
            }
         }
      }
   }
}

typedef struct {
  int starti;
  int startj;
  int startk;
  int TS;
  void* src_addr;
}targs;

void stencil_core(targs *args){
   
   int TS=args->TS;
   int ii,jj,kk;
   int i=args->starti;
   int j=args->startj;
   int k=args->startk;
   for ( ii = 0; ii < TS; ii++) {
      for (jj = 0; jj < TS; jj++) {
         for (kk = 0; kk < TS; kk++) {
            C[i][j][ii][jj] += A[i][k][ii][kk] * B[k][j][kk][jj];
         }
      }
    }

}

int matmul(int N, int M, int P, int TS) 
{
   int NB_I = N / TS;
   int NB_J = P / TS;
   int NB_K = M / TS;
   int i,j,k;
   targs each;
   acemesh_begin_split_task("dgemm_tile");
   for (k = 0; k < NB_K; k++) {
      for (i = 0; i < NB_I; i++) {
         for ( j = 0; j < NB_J; j++) {
            each.TS =TS;
            each.starti=i;
            each.startj=j;
            each.startk=k;
            acemesh_push_rlist(1,&A[i][k][0][0],NORMAL);
            acemesh_push_rlist(1,&B[k][j][0][0],NORMAL);
            acemesh_push_wlist(1,&C[i][j][0][0],NORMAL);
            acemesh_task_generator((TASK_FUNCPTR)stencil_core,(void*)(&each),sizeof(targs));
            acemesh_task_set_type(STENCIL_TASK);
         }
      }
   }
   acemesh_end_split_task();
   acemesh_spawn_and_wait(1);
   return 0;
}

bool verify(int N, int M, int TS, double value) 
{
   printf("begin check!\n");
   int NB_I = N / TS;
   int NB_J = M / TS;
   int i,j,ii,jj;
   bool validates = true;
   for ( i = 0; i < NB_I; i++) {
      for (j = 0; j < NB_J; j++) {
         for (ii = 0; ii < TS; ii++) {
            for ( jj = 0; jj < TS; jj++) {
               if (C[i][j][ii][jj] != value) {
                  validates = false;
                  break;
               }
            }
         }
      }
   }

   return validates;
}

int main(int argc, char* argv[]) 
{
   struct timeval start;
   struct timeval stop;
   double elapsed=0.0;

   if (argc > 6) {
      fprintf(stderr, "input arg error,Usage: %s <N> <M> <P> <BLOCK_SIZE> <NUM_THREADS>\n", argv[0]);
      return 1;
   }
   int num_threads=8;
   int N = 1024;
   int M = 1024;
   int P = 1024;
   int TS = 128;
   if(argc > 1)
      N = atoi(argv[1]);
   if(argc > 2)
      M  = atoi(argv[2]);
   if(argc > 3)
      P = atoi(argv[3]);
   if(argc > 4)
      TS = atoi(argv[4]);
   if(argc > 5)
      num_threads = atoi(argv[5]);
   if (N%TS != 0L) {
      printf("Error: TS (%d) must divide N (%d)\n", TS, N);
      return 1;
   }
   else if (M%TS != 0L) {
      printf("Error: TS (%d) must divide M (%d)\n", TS, M);
      return 1;
   }
   else if (P%TS != 0L) {
      printf("Error: TS (%d) must divide P (%d)\n", TS, P);
      return 1;
   }

   acemesh_runtime_init(num_threads);
   initialize(N, M, P,TS);
   if (A == NULL || B == NULL || C == NULL)
   {
      printf("Error: Can't allocate %ld Bytes\nExiting.\n", (N*M + M*P + N*P)*sizeof(double));
      exit(1);
   }
   gettimeofday(&start, NULL);
   matmul(N, M, P, TS);
   gettimeofday(&stop, NULL);
   acemesh_runtime_shutdown();

   elapsed += (stop.tv_sec - start.tv_sec)+(stop.tv_usec - start.tv_usec)/1000000.0;
   
#ifdef CHECK
   bool validates = verify(N, P, TS, 3.0 + 2.0*1.0*M);
   if (validates)
   { 
      printf("check is passed, correct!\n");
   }else{
      printf("Incorrect result obtained\nExiting.\n");
   }
#endif
   printf("N: %d, M: %d, P: %d, block size: %d, time (s): %lf\n", N, M, P, TS,elapsed);
   return 0;
}
